//! A safe, generic, and extensible hash table implementation in Rust.
//! This implementation is based on the C code provided, converted to Rust
//! with memory safety, type safety, and idiomatic Rust practices.

use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::mem;

/// A trait for types that can be used as keys in the hash table.
/// Implementors must provide a hash function and an equality check.
pub trait HashTableKey: Hash + Eq {
    /// Optional function to free the key when the entry is removed.
    fn free(&self) {}
}

/// A trait for types that can be used as values in the hash table.
/// Implementors must provide a function to free the value when the entry is removed.
pub trait HashTableValue {
    /// Optional function to free the value when the entry is removed.
    fn free(&self) {}
}

/// A struct representing a pair of key and value in the hash table.
#[derive(Debug)]
struct HashTableEntry<K: HashTableKey, V: HashTableValue> {
    key: K,
    value: V,
    next: Option<Box<HashTableEntry<K, V>>>,
}

/// A struct representing the hash table.
pub struct HashTable<K: HashTableKey, V: HashTableValue> {
    table: Vec<Option<Box<HashTableEntry<K, V>>>>,
    table_size: usize,
    entries: usize,
    prime_index: usize,
}

impl<K: HashTableKey, V: HashTableValue> HashTable<K, V> {
    /// Creates a new hash table with the initial size based on the first prime number.
    pub fn new() -> Self {
        let initial_size = hash_table_primes()[0];
        let table = vec![None; initial_size];
        Self {
            table,
            table_size: initial_size,
            entries: 0,
            prime_index: 0,
        }
    }

    /// Frees the memory associated with the hash table, including all entries.
    pub fn free(&mut self) {
        for entry in self.table.iter_mut() {
            if let Some(mut e) = entry.take() {
                e.key.free();
                e.value.free();
                while let Some(mut next) = e.next.take() {
                    next.key.free();
                    next.value.free();
                    e = next;
                }
            }
        }
    }

    /// Inserts a key-value pair into the hash table.
    pub fn insert(&mut self, key: K, value: V) -> bool {
        if (self.entries * 3) / self.table_size > 0 {
            self.enlarge();
        }

        let index = self.hash(&key) % self.table_size;

        if let Some(entry) = self.table[index].as_mut() {
            let mut current = entry;
            loop {
                if current.key == key {
                    current.key.free();
                    current.value.free();
                    current.key = key;
                    current.value = value;
                    return true;
                }
                if let Some(ref mut next) = current.next {
                    current = next;
                } else {
                    break;
                }
            }
            current.next = Some(Box::new(HashTableEntry {
                key,
                value,
                next: None,
            }));
        } else {
            self.table[index] = Some(Box::new(HashTableEntry {
                key,
                value,
                next: None,
            }));
        }

        self.entries += 1;
        true
    }

    /// Looks up a value by key in the hash table.
    pub fn lookup(&self, key: &K) -> Option<&V> {
        let index = self.hash(key) % self.table_size;

        if let Some(entry) = &self.table[index] {
            let mut current = entry;
            loop {
                if current.key == *key {
                    return Some(&current.value);
                }
                if let Some(ref next) = current.next {
                    current = next;
                } else {
                    break;
                }
            }
        }

        None
    }

    /// Removes a key-value pair from the hash table.
    pub fn remove(&mut self, key: &K) -> bool {
        let index = self.hash(key) % self.table_size;

        if let Some(entry) = self.table[index].as_mut() {
            let mut current = entry;
            let mut prev: Option<&mut Box<HashTableEntry<K, V>>> = None;

            loop {
                if current.key == *key {
                    current.key.free();
                    current.value.free();
                    if let Some(p) = prev {
                        p.next = current.next.take();
                    } else {
                        self.table[index] = current.next.take();
                    }
                    self.entries -= 1;
                    return true;
                }
                if let Some(ref mut next) = current.next {
                    prev = Some(current);
                    current = next;
                } else {
                    break;
                }
            }
        }

        false
    }

    /// Returns the number of entries in the hash table.
    pub fn num_entries(&self) -> usize {
        self.entries
    }

    /// Enlarges the hash table to the next prime size.
    fn enlarge(&mut self) {
        let old_table = mem::replace(&mut self.table, vec![None; self.table_size]);
        let old_size = self.table_size;
        self.prime_index += 1;
        self.table_size = hash_table_primes()[self.prime_index];
        self.table = vec![None; self.table_size];

        for mut entry in old_table.into_iter().flatten() {
            while let Some(e) = entry.next.take() {
                self.insert(e.key, e.value);
                entry = e;
            }
            self.insert(entry.key, entry.value);
        }
    }

    /// Computes the hash of a key.
    fn hash(&self, key: &K) -> usize {
        let mut hasher = DefaultHasher::new();
        key.hash(&mut hasher);
        hasher.finish() as usize
    }
}

/// A list of good hash table prime numbers for resizing.
fn hash_table_primes() -> &'static [usize] {
    &[
        193, 389, 769, 1543, 3079, 6151, 12289, 24593, 49157, 98317, 196613, 393241, 786433,
        1572869, 3145739, 6291469, 12582917, 25165843, 50331653, 100663319, 201326611, 402653189,
        805306457, 1610612741,
    ]
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_insert_and_lookup() {
        let mut hash_table = HashTable::new();
        hash_table.insert("key1", 10);
        hash_table.insert("key2", 20);

        assert_eq!(hash_table.lookup(&"key1"), Some(&10));
        assert_eq!(hash_table.lookup(&"key2"), Some(&20));
        assert_eq!(hash_table.lookup(&"key3"), None);
    }

    #[test]
    fn test_remove() {
        let mut hash_table = HashTable::new();
        hash_table.insert("key1", 10);
        hash_table.insert("key2", 20);

        assert!(hash_table.remove(&"key1"));
        assert_eq!(hash_table.lookup(&"key1"), None);
        assert_eq!(hash_table.lookup(&"key2"), Some(&20));
    }

    #[test]
    fn test_num_entries() {
        let mut hash_table = HashTable::new();
        assert_eq!(hash_table.num_entries(), 0);

        hash_table.insert("key1", 10);
        assert_eq!(hash_table.num_entries(), 1);

        hash_table.insert("key2", 20);
        assert_eq!(hash_table.num_entries(), 2);

        hash_table.remove(&"key1");
        assert_eq!(hash_table.num_entries(), 1);
    }
}
